# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: reporter\generators\ds_scan.py
import json, pystache, logging, datetime
from sqlalchemy import select as sql_select
from sqlalchemy import func as sql_func
from sqlalchemy import or_ as sql_or
from sqlalchemy import and_ as sql_and
from sqlalchemy import desc as sql_desc
from reporter.generators.generic import DataSource
from db import Connection
from db.tables.locations import LocationRow, LocationsTable
from db.tables.targets import TargetRow, TargetsTable, TargetTypes
from db.tables.target_vulns import TargetVulnRow, TargetVulnsTable
from db.tables.scan_sessions import ScanSessionRow, ScanSessionsTable
from db.tables.vuln_types import VulnTypeRow, VulnTypesTable
from db.tables.scan_session_vulns import ScanSessionVulnRow, ScanSessionVulnsTable
from db.tables.profiles import ProfileRow, ProfilesTable
from db.tables.users import UsersTable, UserRow
from helpers.address import extract_domain_from_address
from settings_provider import settings
from helpers.constants.jobs import ScanAppTypes

root_location_table_alias = LocationsTable.alias("root_table")
logger = logging.getLogger("reporter.generators.ds_scan")
vuln_type_filters = (
    "cvss3",
    "long_description",
    "name",
    "recommendation",
    "cvss_score",
    "description",
    "cvss2",
    "vt_id",
    "severity",
    "impact",
    "app_id",
)
location_filter = (
    "loc_type",
    "source_id",
    "root_id",
    "parent_id",
    "name",
    "input_data",
    "path",
    "loc_id",
)
vuln_filter = (
    "loc_url",
    "vt_id",
    "vuln_hash",
    "request",
    "source",
    "loc_id",
    "details",
    "app_id",
    "loc_detail",
    "request_secure",
    "request_port",
    "attack_vector",
    "use_ssl",
    "first_seen",
    "status",
)


class DSScanSessionResult(DataSource):
    def __init__(
        self,
        shard_id,
        scan_session_id,
        filter_target_id=None,
        vuln_id_list=None,
        source_type=None,
    ):
        """
        :param shard_id:
        :param scan_session_id:
        :param filter_target_id: result is filtered by target id for multi host scans
        :param vuln_id_list:
        :param source_type:
        """
        super().__init__()
        self.shard_id = shard_id
        self.scan_session_id = scan_session_id
        self.filter_target_id = filter_target_id
        self.vuln_id_list = vuln_id_list
        self.source_type = source_type
        self.scan_info = None
        self.host = None
        self.location_root_filter = None
        self.hide_web_scan_details = False
        self.start_target_id = None
        with Connection(self.shard_id) as (db_connection):
            select_from = ScanSessionsTable.join(
                TargetsTable, TargetRow.target_id == ScanSessionRow.target_id
            )
            select_from = select_from.outerjoin(
                ProfilesTable, ProfileRow.profile_id == ScanSessionRow.used_profile_id
            )
            select_from = select_from.outerjoin(
                UsersTable, UserRow.user_id == TargetRow.owner_id
            )
            s = db_connection.execute(
                sql_select(
                    (
                        ScanSessionRow,
                        TargetRow.address,
                        UserRow.license_type,
                        TargetRow.type.label("target_type"),
                        ProfileRow.name.label("profile_name"),
                    )
                )
                .select_from(select_from)
                .where(ScanSessionRow.scan_session_id == scan_session_id)
            ).fetchone()
            if not s:
                raise Exception("scan or main target not found")
            if s["license_type"] == "trial":
                if s["target_type"] != TargetTypes.DEMO:
                    self.hide_web_scan_details = True
            self.scan_info = dict(s)
            self.start_target_id = self.scan_info["target_id"]
            if self.filter_target_id:
                if self.filter_target_id == self.scan_info["target_id"]:
                    filter_domain = extract_domain_from_address(
                        self.scan_info["address"]
                    )
                    self.host = self.scan_info["address"]
                else:
                    if not self.scan_info["allowed_targets"]:
                        raise Exception("not multiple target scan")
                    for host, filter_target_id in self.scan_info[
                        "allowed_targets"
                    ].items():
                        if self.filter_target_id == filter_target_id:
                            filter_domain = host
                            self.host = host
                            break
                    else:
                        raise Exception("target_id not found %s" % (filter_target_id,))

                q = (
                    sql_select((LocationRow.loc_id, LocationRow.name))
                    .where(LocationRow.scan_session_id == self.scan_session_id)
                    .where(LocationRow.parent_id.is_(None))
                )
                for location in db_connection.execute(q).fetchall():
                    if extract_domain_from_address(location.name) == filter_domain:
                        self.location_root_filter = location.loc_id
                        break
                else:
                    self.location_root_filter = -1

    def info(self):
        """
        required fields: avg_response_time, os, max_response_time, aborted, progress, responsive, duration, technologies
                         host, request_count, aborted_reason, server, sensor, start_date, end_date
        :return:
        """
        from db.tables.scan_session_jobs import ScanSessionJobRow, ScanSessionJobsTable
        from db.tables.scan_session_vulns_stats import (
            ScanSessionVulnsStatsRow,
            ScanSessionVulnsStatsTable,
        )

        select_from = ScanSessionJobsTable.outerjoin(
            ScanSessionVulnsStatsTable,
            ScanSessionVulnsStatsRow.scan_session_id
            == ScanSessionJobRow.scan_session_id,
        )
        select_from = select_from.join(
            ScanSessionsTable,
            ScanSessionRow.scan_session_id == ScanSessionJobRow.scan_session_id,
        )
        scan_session_job_query = (
            sql_select(
                (
                    ScanSessionJobRow,
                    ScanSessionRow.target_id,
                    ScanSessionRow.status.label("session_status"),
                    ScanSessionVulnsStatsRow.vuln_stats,
                )
            )
            .select_from(select_from)
            .where(ScanSessionJobRow.scan_session_id == self.scan_session_id)
        )
        duration = ""
        if self.scan_info["end_date"]:
            pass
        if self.scan_info["start_date"]:
            try:
                duration = self.scan_info["end_date"] - self.scan_info["start_date"]
            except Exception as e:
                logger.exception("error during duration calculus %s", e)

            result = dict(
                host=self.host if self.host else self.scan_info["address"],
                start_date=self.scan_info["start_date"],
                end_date=self.scan_info["end_date"],
                start_url=self.scan_info["address"],
                duration=duration,
                vuln_selection=self.vuln_id_list is not None,
                profile=self.scan_info["profile_name"],
                source_type=self.source_type,
                scan_app_stats={},
            )
            with Connection(self.shard_id) as (db_connection):
                job_stats = db_connection.execute(scan_session_job_query).fetchall()
                if not job_stats:
                    return result
                for job_stat in job_stats:
                    result["scan_app_stats"][job_stat.scanning_app] = job_stats
                    if job_stat.scanning_app not in (
                        ScanAppTypes.WVS,
                        ScanAppTypes.OVAS,
                    ):
                        continue
                    stats = None
                    if job_stat.status in ("processing", "aborting"):
                        stats = settings.get("scan_status_cache").get_entry(
                            self.shard_id,
                            job_stat.scanning_app,
                            self.scan_session_id,
                            enable_db_lookup=False,
                        )
                    if not stats:
                        stats = job_stat.stats
                    if not isinstance(stats, dict):
                        continue
                    main_stats = stats.get("main")
                    if main_stats:
                        if "duration" in main_stats:
                            try:
                                result["duration"] = datetime.timedelta(
                                    seconds=int(main_stats["duration"])
                                )
                            except Exception:
                                pass

                            if "progress" in main_stats:
                                result["progress"] = main_stats["progress"]
                            web_scan_status = main_stats.get("web_scan_status")
                            if web_scan_status:
                                result["avg_response_time"] = web_scan_status.get(
                                    "avg_response_time", 0
                                )
                                result["max_response_time"] = web_scan_status.get(
                                    "avg_response_time", 0
                                )
                                result["locations_found"] = web_scan_status.get(
                                    "locations"
                                )
                                result["request_count"] = web_scan_status.get(
                                    "request_count", 0
                                )
                        host_stats = stats.get("hosts")
                        if job_stat.status != "aborted":
                            if job_stat.scanning_app == ScanAppTypes.WVS:
                                result["responsive"] = False
                        if host_stats:
                            if isinstance(host_stats, dict):
                                for target_id in host_stats.keys():
                                    if self.filter_target_id:
                                        if target_id != self.filter_target_id:
                                            continue
                                        else:
                                            if target_id != self.start_target_id:
                                                continue
                                            host_stat = host_stats[target_id]
                                            result["host"] = host_stat.get("host")
                                            web_scan_status = host_stat.get(
                                                "web_scan_status"
                                            )
                                            if web_scan_status:
                                                result[
                                                    "avg_response_time"
                                                ] = web_scan_status.get(
                                                    "avg_response_time"
                                                )
                                                result[
                                                    "max_response_time"
                                                ] = web_scan_status.get(
                                                    "max_response_time"
                                                )
                                                result[
                                                    "locations_found"
                                                ] = web_scan_status.get("locations")
                                                result[
                                                    "request_count"
                                                ] = web_scan_status.get("request_count")
                                            target_info = host_stat.get("target_info")
                                            if target_info:
                                                result[
                                                    "technologies"
                                                ] = target_info.get("technologies")
                                                result["os"] = target_info.get("os")
                                                result["server"] = target_info.get(
                                                    "server"
                                                )
                                                result[
                                                    "sensor_detected"
                                                ] = target_info.get("sensor_detected")
                                                responsive = target_info.get(
                                                    "responsive"
                                                )
                                                if responsive is True:
                                                    result["responsive"] = True
                                            result["hosts_discovered"] = host_stat.get(
                                                "external_hosts"
                                            )
                                            result["is_starting_host"] = host_stat.get(
                                                "is_starting_host"
                                            )

            return result

    def locations(self):
        """
        required fields: {loc_type, source_id, root_id, parent_id, name, tags, input_data, url, path, loc_id}
        :return:
        """
        if self.location_root_filter == -1:
            return
        select_from = LocationsTable.outerjoin(
            root_location_table_alias,
            sql_and(
                root_location_table_alias.c.loc_id == LocationRow.root_id,
                root_location_table_alias.c.scan_session_id
                == LocationRow.scan_session_id,
            ),
        )
        if self.vuln_id_list:
            select_from = select_from.join(
                ScanSessionVulnsTable,
                sql_and(
                    ScanSessionVulnRow.loc_id == LocationRow.loc_id,
                    ScanSessionVulnRow.scan_session_id == LocationRow.scan_session_id,
                ),
            )
        loc_query = (
            sql_select(
                (LocationRow, root_location_table_alias.c.name.label("root_url"))
            )
            .select_from(select_from)
            .where(LocationRow.scan_session_id == self.scan_session_id)
            .where(LocationRow.path.isnot(None))
            .order_by(LocationRow.path, LocationRow.loc_id)
            .limit(50)
        )
        if self.location_root_filter:
            loc_query = loc_query.where(
                LocationRow.root_id == self.location_root_filter
            )
        if self.vuln_id_list:
            loc_query = loc_query.where(
                ScanSessionVulnRow.vuln_id.in_(self.vuln_id_list)
            )
        cursor = None
        while True:
            locations = []
            if cursor is None:
                q = loc_query
            else:
                q = loc_query.where((LocationRow.path, LocationRow.loc_id) > cursor)
            with Connection(self.shard_id) as (db_connection):
                for location in db_connection.execute(q).fetchall():
                    cursor = (location.path, location.loc_id)
                    if location.source_id is None:
                        logger.error(
                            "incomplete location for loc:id%s", location.loc_id
                        )
                    location = dict(location)
                    root_url = location["root_url"][:-1] if location["root_url"] else ""
                    location["url"] = root_url + location["path"]
                    try:
                        location["input_data"] = json.loads(location["input_data"])[-1]
                    except Exception:
                        location["input_data"] = None

                    locations.append(location)

            if not locations:
                break
            for location in locations:
                yield location

    def vulnerability_types(self):
        """
        required fields: (vuln_type, count)
        vuln_type: cvss3, long_description, name, recommendation, cvss_score, description, cvss, tags, vt_id,
                   references, severity, impact, app_id
        :return:
        """
        if self.location_root_filter == -1:
            return
        select_from = VulnTypesTable.join(
            ScanSessionVulnsTable, ScanSessionVulnRow.vt_id == VulnTypeRow.vt_id
        )
        if self.vuln_id_list is None:
            select_from = select_from.join(
                ScanSessionsTable,
                ScanSessionRow.scan_session_id == ScanSessionVulnRow.scan_session_id,
            )
            select_from = select_from.join(
                TargetVulnsTable,
                sql_and(TargetVulnRow.vuln_hash == ScanSessionVulnRow.vuln_hash),
            )
        if self.location_root_filter:
            select_from = select_from.join(
                LocationsTable,
                sql_and(
                    LocationRow.loc_id == ScanSessionVulnRow.loc_id,
                    LocationRow.scan_session_id == ScanSessionVulnRow.scan_session_id,
                ),
            )
        vt_query = sql_select(
            (VulnTypeRow, sql_func.count(VulnTypeRow.vt_id).label("vuln_count"))
        ).select_from(select_from)
        if self.vuln_id_list is None:
            vt_query = vt_query.where(TargetVulnRow.status == "open")
        if self.location_root_filter:
            vt_query = vt_query.where(LocationRow.root_id == self.location_root_filter)
        vt_query = (
            vt_query.where(ScanSessionVulnRow.scan_session_id == self.scan_session_id)
            .group_by(VulnTypeRow.vt_id)
            .order_by(
                sql_desc(VulnTypeRow.severity), VulnTypeRow.name, VulnTypeRow.vt_id
            )
            .limit(50)
        )
        if self.vuln_id_list:
            vt_query = vt_query.where(ScanSessionVulnRow.vuln_id.in_(self.vuln_id_list))
        cursor = None
        while True:
            if cursor is None:
                q = vt_query
            else:
                q = vt_query.where(
                    sql_or(
                        VulnTypeRow.severity < cursor[0],
                        sql_and(
                            VulnTypeRow.severity == cursor[0],
                            sql_or(
                                VulnTypeRow.name > cursor[1],
                                sql_and(
                                    VulnTypeRow.name == cursor[1],
                                    VulnTypeRow.vt_id > cursor[2],
                                ),
                            ),
                        ),
                    )
                )
            vtypes = []
            with Connection(self.shard_id) as (db_connection):
                for vt in db_connection.execute(q).fetchall():
                    cursor = (vt.severity, vt.name, vt.vt_id)
                    vt = dict(vt)
                    refs = []
                    if vt["refs"]:
                        for ref in vt["refs"]:
                            refs.append(dict(name=ref[0], url=ref[1]))

                    vt["refs"] = refs
                    vt["type"] = self.extract_vuln_type_from_tags(vt["tags"])
                    vtypes.append((vt, vt["vuln_count"]))

                if not vtypes:
                    break
                for vtype in vtypes:
                    yield vtype

    def vulnerabilities(self, vtype=None, location=None, tags=None, severity=None):
        """
        tags is not exact match filter, is overlap filter !!!
        
        required fields: (vuln, location, vuln_type)
        vuln: loc_url, host, vt_id, vuln_hash, request, source, loc_id, details, app_id, loc_detail
        location: loc_type, source_id, root_id, parent_id, name, tags, input_data, url, path, loc_id
        vuln_type: cvss3, long_description, name, recommendation, cvss_score, description, cvss, tags, vt_id,
                   references, severity, impact, app_id
        :param vtype:
        :param location:
        :param tags:
        :param severity:
        :return:
        """
        if self.location_root_filter == -1:
            return
        if tags:
            tags = list(tags)
        select_from = ScanSessionVulnsTable.join(
            VulnTypesTable, VulnTypeRow.vt_id == ScanSessionVulnRow.vt_id
        )
        select_from = select_from.join(
            LocationsTable,
            sql_and(
                LocationRow.loc_id == ScanSessionVulnRow.loc_id,
                LocationRow.scan_session_id == ScanSessionVulnRow.scan_session_id,
            ),
        )
        select_from = select_from.outerjoin(
            root_location_table_alias,
            sql_and(
                root_location_table_alias.c.loc_id == LocationRow.root_id,
                root_location_table_alias.c.scan_session_id
                == LocationRow.scan_session_id,
            ),
        )
        if self.vuln_id_list is None:
            select_from = select_from.join(
                ScanSessionsTable,
                ScanSessionRow.scan_session_id == ScanSessionVulnRow.scan_session_id,
            )
            select_from = select_from.join(
                TargetVulnsTable,
                sql_and(TargetVulnRow.vuln_hash == ScanSessionVulnRow.vuln_hash),
            )
        query = (
            sql_select(
                (
                    ScanSessionVulnRow,
                    VulnTypeRow.severity,
                    VulnTypeRow.name,
                    VulnTypeRow.details_template,
                    VulnTypeRow.app_id,
                    LocationRow.path,
                    LocationRow.loc_type,
                    LocationRow.source_id,
                    LocationRow.root_id,
                    LocationRow.parent_id,
                    LocationRow.input_data,
                    LocationRow.tags.label("loc_tags"),
                    root_location_table_alias.c.name.label("root_url"),
                    VulnTypeRow.tags.label("vt_tags"),
                    VulnTypeRow.cvss2,
                    VulnTypeRow.cvss3,
                    VulnTypeRow.cvss_score,
                    VulnTypeRow.impact,
                    VulnTypeRow.description,
                    VulnTypeRow.recommendation,
                    VulnTypeRow.long_description,
                    VulnTypeRow.refs,
                    TargetVulnRow.first_seen,
                    TargetVulnRow.status,
                )
            )
            .select_from(select_from)
            .where(ScanSessionVulnRow.scan_session_id == self.scan_session_id)
        )
        if self.vuln_id_list:
            query = query.where(ScanSessionVulnRow.vuln_id.in_(self.vuln_id_list))
        else:
            query = query.where(TargetVulnRow.status == "open")
        if self.location_root_filter:
            query = query.where(LocationRow.root_id == self.location_root_filter)
        if vtype:
            query = query.where(ScanSessionVulnRow.vt_id == vtype["vt_id"])
        if location:
            query = query.where(ScanSessionVulnRow.loc_id == location["loc_id"])
        if severity is not None:
            query = query.where(VulnTypeRow.severity == severity)
        if tags:
            query = query.where(VulnTypeRow.tags.overlap(tags))
        query = query.order_by(
            sql_desc(VulnTypeRow.severity),
            VulnTypeRow.name,
            LocationRow.path,
            ScanSessionVulnRow.loc_detail,
            ScanSessionVulnRow.vuln_id,
        ).where(LocationRow.path.isnot(None))
        query = query.limit(50)
        cursor = None
        while True:
            if cursor is None:
                q = query
            else:
                q = query.where(
                    sql_or(
                        VulnTypeRow.severity < cursor[0],
                        sql_and(
                            VulnTypeRow.severity == cursor[0],
                            sql_or(
                                VulnTypeRow.name > cursor[1],
                                sql_and(
                                    VulnTypeRow.name == cursor[1],
                                    sql_or(
                                        LocationRow.path > cursor[2],
                                        sql_and(
                                            LocationRow.path == cursor[2],
                                            sql_or(
                                                ScanSessionVulnRow.loc_detail
                                                > cursor[3],
                                                sql_and(
                                                    ScanSessionVulnRow.loc_detail
                                                    == cursor[3],
                                                    ScanSessionVulnRow.vuln_id
                                                    > cursor[4],
                                                ),
                                            ),
                                        ),
                                    ),
                                ),
                            ),
                        ),
                    )
                )
            vulns = []
            with Connection(self.shard_id) as (db_connection):
                for vuln in db_connection.execute(q).fetchall():
                    cursor = (
                        vuln.severity,
                        vuln.name,
                        vuln.path,
                        vuln.loc_detail,
                        vuln.vuln_id,
                    )
                    if vuln.source_id is None:
                        logger.error(
                            "incomplete location for vuln_id:%s, loc_id:%s",
                            vuln.vuln_id,
                            vuln.loc_id,
                        )
                        continue
                    data = dict(vuln)
                    root_url = data["root_url"][:-1] if data["root_url"] else ""
                    data["loc_url"] = root_url + data["path"]
                    try:
                        if data["details_type"] == "json":
                            data["details"] = pystache.render(
                                data["details_template"], json.loads(data["details"])
                            )
                    except Exception as e:
                        logger.exception(
                            "render_details failed with %s for %s" % (e, vuln)
                        )

                    vuln = {key: data[key] for key in vuln_filter}
                    vuln["host"] = data["root_url"]
                    location = {key: data[key] for key in location_filter}
                    location["url"] = data["loc_url"]
                    location["tags"] = data["loc_tags"]
                    vuln_type = {key: data[key] for key in vuln_type_filters}
                    vuln_type["tags"] = data["vt_tags"]
                    refs = []
                    if data["refs"]:
                        for ref in data["refs"]:
                            refs.append(dict(name=ref[0], url=ref[1]))

                    vuln_type["references"] = refs
                    vuln_type["type"] = self.extract_vuln_type_from_tags(
                        vuln_type["tags"]
                    )
                    vulns.append((vuln, location, vuln_type))

            if not vulns:
                break
            for vuln in vulns:
                yield vuln
